from volprims_linear import *
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
import torch.optim as optim
import os
import slangpy as spy
import pathlib
from volprims_linear import PinHoleCamera
from dataset import VolumetricPrimitivesDataset
import random

# MAX_NUM_INSTANCES = 1024 * 1024
MAX_NUM_INSTANCES = 8 * 1024 * 1024 # a maximum of 8M instances
NUM_PRIMITIVES = 4096
IMAGE_WIDTH = 1024
IMAGE_HEIGHT = 1024


device = spy.create_device(
    type=spy.DeviceType.vulkan, # dx12 does not have support for fp32 atomic accumulation
    include_paths=[
        pathlib.Path(__file__).parent.absolute(),
    ],
    enable_debug_layers=True,
    enable_cuda_interop=True,
    enable_print=True)

# torch.autograd.set_detect_anomaly(True)

def load_target_image(image_path, image_width=256, image_height=256):
    '''
    Load the target image and convert it to a torch.Tensor
    @param image_path: Path to the target image
    @param image_size: Size to which the image will be resized (default is 256)
    @return: A numpy array of the image with shape (image_size, image_size, 4)
    '''
    if os.path.exists(image_path):
        image = Image.open(image_path)
        image = image.resize((image_width, image_height), Image.Resampling.LANCZOS)
        image_np = np.array(image) / 255.0
        
        if image_np.shape[2] == 3:
            alpha = np.ones((image_width, image_height, 1))
            image_np = np.concatenate([image_np, alpha], axis=2)
            
        return image_np
    else:
        print(f"Target image {image_path} does not exist.")
        return np.zeros((image_width, image_height, 4), dtype=np.float32)

def init_primitives(num_primitives=256):
    '''
    Initialize the volume primitives with random parameters.
    @param num_primitives: Number of primitives to initialize
    @return: An instance of UniformVolumePrimitives with initialized parameters
    '''
    primitives = UniformVolumePrimitives(num_spheres=num_primitives)
    primitives.random_initialization(device=default_device)
    
    return primitives

def setup_camera():
    """
    Initialize the camera for rendering.
    The camera is positioned to look down at the origin from above, with a 60-degree
    field of view. The camera's up vector is tilted to ensure a 45-degree downward pitch.
    @return: An instance of PinHoleCamera configured for the scene
    """
    cam_pos_y = -world_radius * 1.5
    cam_pos_z = world_radius * 1.5
    
    camera_position = torch.tensor([0.0, cam_pos_y, cam_pos_z], device=default_device)
    look_at_point = torch.tensor([0.0, 0.0, 0.0], device=default_device)
    up_vector = torch.tensor([0.0, 1.0, 1.0], device=default_device) 
    fov_degrees = 60.0

    camera = PinHoleCamera(
        position=camera_position,
        look_at=look_at_point,
        up=up_vector,
        fov=fov_degrees,
        device=default_device
    )
    
    return camera

def make_camera_vectors(camera : PinHoleCamera, aspect_ratio=1.0):
    """
    Create the up and right vectors for the camera.
    ndc.x * right + ndc.y * up + ndc.z * forward -> an unnormalized direction vector
    pointing to the z=1 plane in the camera's coordinate system.
    @param camera: The camera instance
    @return: A tuple containing the numpy up vector and right vector of the camera
    """
    up_vec = camera.up
    direction = camera.get_direction()
    right_vec = torch.cross(direction , up_vec)
    up_vec = torch.cross(right_vec, direction)
    right_vec = right_vec / torch.norm(right_vec)
    up_vec = up_vec / torch.norm(up_vec)

    tan_fov_2 = np.tan(np.radians(camera.fov) / 2)
    up_vec = up_vec * tan_fov_2
    right_vec = right_vec * tan_fov_2 * aspect_ratio
    return up_vec.detach().cpu().numpy(), right_vec.detach().cpu().numpy()

class DeviceUniformVolumePrimitives:
    def __init__(self, num_primitives : int):
        self.blob_float_count = 0
        self.blob = device.create_buffer(
            size = num_primitives * 14 * 4, # 14 floats per sphere 
            usage = spy.BufferUsage.unordered_access 
            | spy.BufferUsage.copy_source 
            | spy.BufferUsage.copy_destination,
        )
        self.d_blob = device.create_buffer(
            size = num_primitives * 14 * 4, # 14 floats per sphere 
            usage = spy.BufferUsage.unordered_access 
            | spy.BufferUsage.copy_source 
            | spy.BufferUsage.copy_destination,
        )
        self.blob_float_count = num_primitives * 14

def upload_primitives(primitives : UniformVolumePrimitives):
    """
    Upload the volume primitives to the device.
    @param primitives: The volume primitives instance
    @return DeviceUniformPrimitives: The device data for the primitives
    """
    # Upload the rotations, scales, opacities, positions and colors to the device
    numpy_packed_array = np.zeros((len(primitives.centers), 14), dtype=np.float32)
    numpy_packed_array[:, 0:4] = primitives.rotations.detach().cpu().numpy()  # rotations
    numpy_packed_array[:, 4:7] = primitives.scales.detach().cpu().numpy()     # scales
    numpy_packed_array[:, 7:8] = primitives.opacities.detach().cpu().numpy()  # opacities
    numpy_packed_array[:, 8:11] = primitives.centers.detach().cpu().numpy()   # centers
    numpy_packed_array[:, 11:14] = primitives.colors.detach().cpu().numpy()   # colors
    numpy_packed_array = numpy_packed_array.flatten()
    device_data = DeviceUniformVolumePrimitives(num_primitives=primitives.num_spheres)
    device_data.blob.copy_from_numpy(numpy_packed_array)
    return device_data

def setup_uniforms(camera : PinHoleCamera,
                   primitives : UniformVolumePrimitives, 
                   dev_primitives : DeviceUniformVolumePrimitives,
                   image_width=256, image_height=256):
    """
    Set up the uniforms for rendering.
    @param camera: The camera instance
    @param primitives: The volume primitives instance
    @param image_width: Width of the rendered image
    @param image_height: Height of the rendered image
    @return: A dictionary containing the uniforms for rendering
    """
    up_vec, right_vec = make_camera_vectors(camera, aspect_ratio=image_width / image_height)
    return {
        "camera_position": camera.position.detach().cpu().numpy(),
        "camera_fov_y": 2 * np.tan(np.radians(camera.fov) / 2),
        "camera_direction": camera.get_direction().detach().cpu().numpy(),
        "camera_aspect_ratio": image_width / image_height,
        "film_dimensions": (image_width, image_height),
        "blob_float_count": dev_primitives.blob_float_count,
        "primitive_count": primitives.num_spheres,
        "camera_right": right_vec,
        "max_num_instances": MAX_NUM_INSTANCES,
        "camera_up": up_vec,
        "render_tile_pixel_index_offset": 0,
        "camera_view": camera.get_view_matrix(rhs = False).detach().cpu().numpy(),
        "camera_perspective_view": camera.get_perspective_view_matrix(rhs = False).detach().cpu().numpy(),
        "tile_dimensions": ((image_width + 15) // 16, (image_height + 15) // 16),
        "padding1": (0, 0)
    }

def update_uniforms_for_camera(camera : PinHoleCamera, uniforms : dict):
    """
    Update the uniforms for rendering based on the camera parameters.
    @param camera: The camera instance
    @param uniforms: The uniforms dictionary to update
    """
    up_vec, right_vec = make_camera_vectors(camera, aspect_ratio=uniforms["camera_aspect_ratio"])
    uniforms["camera_position"] = camera.position.detach().cpu().numpy()
    uniforms["camera_fov_y"] = 2 * np.tan(np.radians(camera.fov) / 2)
    uniforms["camera_direction"] = camera.get_direction().detach().cpu().numpy()
    uniforms["camera_right"] = right_vec
    uniforms["camera_up"] = up_vec
    uniforms["camera_view"] = camera.get_view_matrix(rhs=False).detach().cpu().numpy()
    uniforms["camera_perspective_view"] = camera.get_perspective_view_matrix(rhs=False).detach().cpu().numpy()

class SlangProgramResources:
    def __init__(self, num_primitives, image_width, image_height):
        self.active_primitive_list_buf = device.create_buffer(
            size=num_primitives * 4,  # 4 bytes per primitive index
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label='active_primitive_list_buf'
        )
        self.active_primitive_count =  device.create_buffer(
            size=4,  # 4 bytes for the count
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label = 'active_primitive_count'
        )
        self.active_instance_count = device.create_buffer(
            size=4,  # 4 bytes for the count
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label = 'active_instance_count'
        )
        self.active_instance_sorting_key_buf = device.create_buffer(
            size=MAX_NUM_INSTANCES * 4,  # 4 bytes per key
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label = 'active_instance_sorting_key_buf'
        )
        self.active_instance_primitive_list_buf = device.create_buffer(
            size=MAX_NUM_INSTANCES * 4,  # 4 bytes per instance index
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label= 'active_instance_primitive_list_buf'
        )
        self.active_instance_sorting_key_sorted_buf = device.create_buffer(
            size=MAX_NUM_INSTANCES * 4,  # 4 bytes per key
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label= 'active_instance_sorting_key_sorted_buf'
        )
        self.active_instance_primitive_list_sorted_buf = device.create_buffer(
            size=MAX_NUM_INSTANCES * 4,  # 4 bytes per index
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label= 'active_instance_primitive_list_sorted_buf'
        )
        self.radix_sort_temp_keys_buf = device.create_buffer(
            size=MAX_NUM_INSTANCES * 4,  # 4 bytes per key
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label= 'radix_sort_temp_keys_buf'
        )
        self.radix_sort_temp_values_buf = device.create_buffer(
            size=MAX_NUM_INSTANCES * 4,  # 4 bytes per value
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label= 'radix_sort_temp_values_buf'
        )
        ELEMENTS_PER_SEGMENT = 1024
        max_segment_count =  (MAX_NUM_INSTANCES + ELEMENTS_PER_SEGMENT - 1) // ELEMENTS_PER_SEGMENT
        BINS_PER_PASS = 256
        self.radix_sort_bins_buf = device.create_buffer(
            size=max_segment_count * BINS_PER_PASS * 4,  # 4 bytes per bin
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label= 'radix_sort_bins_buf'
        )
        self.radix_sort_sum_bins_buf = device.create_buffer(
            size=BINS_PER_PASS * 4,
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label= 'radix_sort_sum_bins_buf'
        )
        self.indirect_dispatch_args_buf = device.create_buffer(
            size=4 * 4,  # 4 uint32_t for indirect dispatch args
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.indirect_argument | spy.BufferUsage.unordered_access,
            label= 'indirect_dispatch_args_buf'
        )
        self.prims_indirect_dispatch_args_buf = device.create_buffer(
            size=4 * 4,  # 4 uint32_t for indirect dispatch args
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.indirect_argument | spy.BufferUsage.unordered_access,
            label= 'prims_indirect_dispatch_args_buf'
        )
        num_tiles = ((image_width + 15) // 16) * ((image_height + 15) // 16)
        self.tile_instance_primitive_list_offsets_buf = device.create_buffer(
            size= num_tiles * 4,  # 4 bytes per offset
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label= 'tile_instance_primitive_list_offsets_buf'
        )
        self.tile_instance_primitive_list_counts_buf = device.create_buffer(
            size= num_tiles * 4,  # 4 bytes per count
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label= 'tile_instance_primitive_list_counts_buf'
        )
        self.forward_state_cache_buf = device.create_buffer(
            size=(MAX_NUM_INSTANCES + num_tiles) * 6 * 4 * 32,
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label= 'forward_state_cache_buf'
        )
        self.adam_first_moment = device.create_buffer(
            size=num_primitives * 14 * 4,  # Assuming 14 floats per primitive
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label= 'adam_first_moment_buf'
        )
        self.adam_second_moment = device.create_buffer(
            size=num_primitives * 14 * 4,  # Assuming 14 floats per primitive
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            label= 'adam_second_moment_buf'
        )
        self.total_loss = device.create_buffer(
            size = 4,
            memory_type=spy.MemoryType.device_local,
            usage=spy.BufferUsage.unordered_access,
            format=spy.Format.r32_float,
            label= 'total_loss'
        )
        self.output_texture = device.create_texture(
            format = spy.Format.rgba32_float,
            width = image_width,
            height = image_height,
            usage = spy.TextureUsage.unordered_access | spy.TextureUsage.shader_resource,
            label= 'output_texture'
        )
        self.d_output_texture = device.create_texture(
            format = spy.Format.rgba32_float,
            width = image_width,
            height = image_height,
            usage = spy.TextureUsage.unordered_access | spy.TextureUsage.shader_resource,
            label= 'd_output_texture'
        )
        self.target_texture = device.create_texture(
            format = spy.Format.rgba32_float,
            width = image_width,
            height = image_height,
            usage = spy.TextureUsage.shader_resource,
            label= 'target_texture'
        )
        self.debug_buf = device.create_buffer(
            size=16 * 1024,
            format=spy.Format.r32_float,
            memory_type=spy.MemoryType.read_back,
            usage=spy.BufferUsage.copy_destination | spy.BufferUsage.unordered_access,
            label= 'debug_buf'
        )
    def bind(self, pass_encoder : spy.ComputePassEncoder, pipeline : spy.ComputePipeline, ub_dict : dict):
        shader_obj =  pass_encoder.bind_pipeline(pipeline)
        cursor = spy.ShaderCursor(shader_obj)
        cursor.ub = ub_dict
        cursor.active_primitive_list_buf = self.active_primitive_list_buf
        cursor.active_primitive_count = self.active_primitive_count
        cursor.active_instance_sorting_key_buf = self.active_instance_sorting_key_buf
        cursor.active_instance_primitive_list_buf = self.active_instance_primitive_list_buf
        cursor.active_instance_sorting_key_sorted_buf = self.active_instance_sorting_key_sorted_buf
        cursor.active_instance_primitive_list_sorted_buf = self.active_instance_primitive_list_sorted_buf
        cursor.active_instance_count = self.active_instance_count
        cursor.tile_instance_primitive_list_offsets_buf = self.tile_instance_primitive_list_offsets_buf
        cursor.tile_instance_primitive_list_counts_buf = self.tile_instance_primitive_list_counts_buf
        cursor.forward_state_cache_buf = self.forward_state_cache_buf
        cursor.adam_first_moment = self.adam_first_moment
        cursor.adam_second_moment = self.adam_second_moment
        cursor.total_loss = self.total_loss;
        cursor.output_texture = self.output_texture
        cursor.target_texture = self.target_texture
        cursor.d_output_texture = self.d_output_texture
        cursor.debug_buf = self.debug_buf
        return cursor

class UniformVolumePrimitivesOptimizer:

    def reinit_primitives(self):
        self.primitives = init_primitives(self.num_primitives)
        
        # Fix them to 0 for testing purposes
        # self.primitives.centers[:, :] = -1
        # self.primitives.scales[:, :] = 4
        # print(self.primitives.centers)
        # print(self.primitives.scales)
        # print(self.primitives.rotations)
        # self.primitives.centers[0] = torch.from_numpy(np.array([-2.7774, -8.5170,  9.1381])).cuda()
        # self.primitives.scales_raw[0] = torch.from_numpy(np.array([0.0527, 0.7586, 0.5481])).cuda()
        # self.primitives.rotations_raw[0] = torch.from_numpy(np.array([ 0.0421, -0.8558, -1.8474,  0.3509])).cuda()

        # print(self.primitives.centers[0])
        # print(self.primitives.scales_raw[0])
        # print(self.primitives.rotations_raw[0])
        self.dev_primitives = upload_primitives(self.primitives)

    def __init__(self):
        self.num_primitives = NUM_PRIMITIVES
        self.image_width = IMAGE_WIDTH
        self.image_height = IMAGE_HEIGHT
        self.reinit_primitives()
        self.camera = setup_camera()
        self.uniforms = setup_uniforms(self.camera, self.primitives, self.dev_primitives,
                                        self.image_width, self.image_height)
        self.resources = SlangProgramResources(self.num_primitives, self.image_width,
                                            self.image_height)
        self.dataset = VolumetricPrimitivesDataset("volprims_dataset_alpha", image_width=self.image_width, image_height=self.image_height)
        # target = load_target_image("sample.png", self.image_width, self.image_height)
        # self.resources.target_texture.copy_from_numpy(target.astype(np.float32))

        
        '''
            "ClearCounters",
            "FilterActivePrimitives",
            "PrecomputeActivePrimitives",
            "RadixSortScan",
            "RadixSortSum",
            "RadixSortSumBins",
            "RadixSortScatter",
            "FindTileBoundaries",
            "CountTileInstances",
            "Render",
            "Render_BackwardPropagation",
            "UpdateParameters"
        '''
        self.ClearCounters_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=[
            "ClearCounters"
            ]
        )
        self.ClearCounters = device.create_compute_kernel(self.ClearCounters_program)

        self.FilterActivePrimitives_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=[
            "FilterActivePrimitives"
            ]
        )
        self.FilterActivePrimitives = device.create_compute_kernel(self.FilterActivePrimitives_program)

        self.PrecomputeActivePrimitives_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=[
            "PrecomputeActivePrimitives"
            ]
        )
        self.PrecomputeActivePrimitives = device.create_compute_kernel(self.PrecomputeActivePrimitives_program)

        self.EncodeDispatchArgs_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=[
                    "EncodeDispatchArgs"
            ]
        )
        self.EncodeDispatchArgs = device.create_compute_kernel(self.EncodeDispatchArgs_program)

        self.RadixSortScan_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=[
            "RadixSortScan"
            ]
        )
        self.RadixSortScan = device.create_compute_kernel(self.RadixSortScan_program)

        self.RadixSortSum_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=[
            "RadixSortSum"
            ]
        )
        self.RadixSortSum = device.create_compute_kernel(self.RadixSortSum_program)

        self.RadixSortSumBins_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=[
            "RadixSortSumBins"
            ]
        )
        self.RadixSortSumBins = device.create_compute_kernel(self.RadixSortSumBins_program)

        self.RadixSortScatter_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=[
            "RadixSortScatter"
            ]
        )
        self.RadixSortScatter = device.create_compute_kernel(self.RadixSortScatter_program)

        self.FindTileBoundaries_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=[
            "FindTileBoundaries"
            ]
        )
        self.FindTileBoundaries = device.create_compute_kernel(self.FindTileBoundaries_program)

        self.CountTileInstances_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=[
            "CountTileInstances"
            ]
        )
        self.CountTileInstances = device.create_compute_kernel(self.CountTileInstances_program)

        self.Render_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=[
            "Render"
            ]
        )
        self.Render = device.create_compute_kernel(self.Render_program)

        self.Render_BackwardPropagation_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=[
            "Render_BackwardPropagation"
            ]
        )
        self.Render_BackwardPropagation = device.create_compute_kernel(self.Render_BackwardPropagation_program)

        self.UpdateParameters_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=[
            "UpdateParameters"
            ]
        )
        self.UpdateParameters = device.create_compute_kernel(self.UpdateParameters_program)

        self.RegularizeParameters_program = device.load_program(
            module_name="slang/optimizer.slang",
            entry_point_names=["RegularizeParameters"]
        )
        self.RegularizeParameters = device.create_compute_kernel(self.RegularizeParameters_program)

        self.window = spy.Window(width=self.image_width, height=self.image_height,
                             title="Optimizer", resizable=False)
        self.surface = device.create_surface(self.window)
        self.surface.configure(
            format=spy.Format.rgba8_unorm_srgb,
            width=self.window.width,
            height=self.window.height,
            vsync=False
        )

        self.window.on_keyboard_event = self.on_keyboard_event
        self.window.on_mouse_event = self.on_mouse_event
        
        self.ui = spy.ui.Context(device)
        self.setup_ui()
    
    def main(self):
        frame = 0
        timer = spy.Timer()

        random_index = 0

        loss_curve = []

        max_num_frames = 20000
        
        while not self.window.should_close() and frame < max_num_frames:
            dt = timer.elapsed_s()

            self.fps_text.text = f"FPS: {1.0 / dt:.2f} | Frame: {frame}"

            timer.reset()

            self.window.process_events()
            self.ui.process_events()

            surface_texture = self.surface.acquire_next_image()
            if not surface_texture:
                continue

            if not self.paused.value:
                random_index = random.randint(0, len(self.dataset) - 1)
            target, train_camera = self.dataset[random_index]
            update_uniforms_for_camera(train_camera, self.uniforms)
            self.resources.target_texture.copy_from_numpy(target.astype(np.float32))


            command_encoder = device.create_command_encoder()

            self.render(command_encoder)
            
            command_encoder.blit(surface_texture, self.resources.output_texture)

            self.ui.new_frame(surface_texture.width, surface_texture.height)
            self.ui.render(surface_texture, command_encoder)

            device.submit_command_buffer(command_encoder.finish())

            # device.wait_for_idle()
            # readback = self.resources.debug_buf.to_numpy()
            # print(readback[:4])
            loss_value = self.resources.total_loss.to_numpy()[0]
            loss_lerp = 0.005
            if self.smoothed_loss_value == 0: loss_lerp = 1.0
            self.smoothed_loss_value = loss_lerp * loss_value + (1 - loss_lerp) * self.smoothed_loss_value
            self.loss_text.text = f"SmoothedLoss: {self.smoothed_loss_value:.6f} | Random Index: {random_index}"

            if frame % 200 == 0: loss_curve.append((frame, self.smoothed_loss_value))

            del surface_texture

            self.surface.present()

            frame += 1
        
        device.wait()

        # Plot loss curve and save to file
        if len(loss_curve) > 0:
            frames, losses = zip(*loss_curve)
            plt.figure(figsize=(10, 5))
            plt.plot(frames, losses, label='Loss Curve')
            plt.xlabel('Frame')
            plt.ylabel('Loss')
            plt.title(f"Loss (prims: {self.num_primitives})")
            plt.legend()
            plt.grid(True)
            plt.savefig(f"loss_prims_{self.num_primitives}.png")
            plt.show()

            # also save the loss curve to a text file
            with open(f"loss_prims_{self.num_primitives}.txt", "w") as f:
                for frame, loss in loss_curve:
                    f.write(f"{frame} {loss}\n")

    def encode_dispatch(self, command_encoder: spy.CommandEncoder,
                         count_buffer : spy.Buffer, out_command : spy.Buffer,
                           threads_per_group : int):
        encode_ub = {
            "NumThreadsPerGroup" : threads_per_group,
            "Padding": (0, 0, 0)
        }
        self.EncodeDispatchArgs.dispatch((1, 1, 1), vars = {
            "UB_Dispatch": encode_ub,
            "DispatchArgsNumThreads": count_buffer,
            "RWDispatchArgs": out_command
        }, command_encoder=command_encoder)

    def radix_sort(self, command_encoder: spy.CommandEncoder, src_keys, src_values, dst_keys, dst_values, count):
        
        ELEMENTS_PER_SEGMENT = 1024
        BINS_PER_PASS = 256
        self.encode_dispatch(command_encoder, count, self.resources.indirect_dispatch_args_buf, ELEMENTS_PER_SEGMENT)

        curr_keys = src_keys
        curr_values = src_values
        curr_dst_keys = self.resources.radix_sort_temp_keys_buf
        curr_dst_values = self.resources.radix_sort_temp_values_buf
        for i in range(4):
            bit_shift = i * 8
            sort_ub = {
                "NumElements": 0,
                "BitShift": bit_shift,
                "Padding0": 0,
                "Padding1": 0
            }
            with command_encoder.begin_compute_pass() as pass_encoder:
                shader_obj = pass_encoder.bind_pipeline(self.RadixSortScan.pipeline)
                cursor = spy.ShaderCursor(shader_obj)
                cursor.Keys   = curr_keys
                cursor.Values = curr_values
                cursor.RWBins = self.resources.radix_sort_bins_buf
                cursor.Count  = count
                cursor.UB_RadixSort = sort_ub
                pass_encoder.dispatch_compute_indirect(self.resources.indirect_dispatch_args_buf)
            with command_encoder.begin_compute_pass() as pass_encoder:
                shader_obj = pass_encoder.bind_pipeline(self.RadixSortSum.pipeline)
                cursor = spy.ShaderCursor(shader_obj)
                cursor.RWBins = self.resources.radix_sort_bins_buf
                cursor.Count  = count
                cursor.UB_RadixSort = sort_ub
                pass_encoder.dispatch_compute((BINS_PER_PASS, 1, 1))
            with command_encoder.begin_compute_pass() as pass_encoder:
                shader_obj = pass_encoder.bind_pipeline(self.RadixSortSumBins.pipeline)
                cursor = spy.ShaderCursor(shader_obj)
                cursor.Bins = self.resources.radix_sort_bins_buf
                cursor.RWSumBins = self.resources.radix_sort_sum_bins_buf
                cursor.Count = count
                cursor.UB_RadixSort = sort_ub
                pass_encoder.dispatch_compute((1, 1, 1))
            with command_encoder.begin_compute_pass() as pass_encoder:
                shader_obj = pass_encoder.bind_pipeline(self.RadixSortScatter.pipeline)
                cursor = spy.ShaderCursor(shader_obj)
                cursor.Keys = curr_keys
                cursor.Values = curr_values
                cursor.RWOutKeys = curr_dst_keys
                cursor.RWOutValues = curr_dst_values
                cursor.Bins = self.resources.radix_sort_bins_buf
                cursor.SumBins = self.resources.radix_sort_sum_bins_buf
                cursor.Count = count
                cursor.UB_RadixSort = sort_ub
                pass_encoder.dispatch_compute_indirect(self.resources.indirect_dispatch_args_buf)
            if i != 0:
                # Swap the buffers for the next iteration
                temp_keys, temp_values = curr_dst_keys, curr_dst_values
                curr_dst_keys, curr_dst_values = curr_keys, curr_values
                curr_keys, curr_values = temp_keys, temp_values
            else:
                curr_keys, curr_values = curr_dst_keys, curr_dst_values
                curr_dst_keys, curr_dst_values = dst_keys, dst_values



    def render(self, command_encoder: spy.CommandEncoder):
        """
        Render the scene using the volume primitives and camera.
        @param command_encoder: The command encoder for issuing rendering commands
        """
        COMPUTE_DEFAULT_GROUP_SIZE = 64
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("ClearCounters", (0, 0, 0))
            cursor = self.resources.bind(pass_encoder, self.ClearCounters.pipeline, self.uniforms)
            pass_encoder.dispatch_compute((1, 1, 1))
            pass_encoder.pop_debug_group()
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("FilterActivePrimitives", (1, 0, 0))
            cursor = self.resources.bind(pass_encoder, self.FilterActivePrimitives.pipeline, self.uniforms)
            cursor.primitive_data_buf = self.dev_primitives.blob
            pass_encoder.dispatch_compute((
                (self.primitives.num_spheres + COMPUTE_DEFAULT_GROUP_SIZE - 1) 
                // COMPUTE_DEFAULT_GROUP_SIZE, 1, 1))
            pass_encoder.pop_debug_group()
        self.encode_dispatch(command_encoder, self.resources.active_primitive_count,
                             self.resources.prims_indirect_dispatch_args_buf, COMPUTE_DEFAULT_GROUP_SIZE)
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("Precompute Active Primitives", (0, 1, 0))
            cursor = self.resources.bind(pass_encoder, self.PrecomputeActivePrimitives.pipeline, self.uniforms)
            cursor.primitive_data_buf = self.dev_primitives.blob
            pass_encoder.dispatch_compute_indirect(self.resources.prims_indirect_dispatch_args_buf)
            pass_encoder.pop_debug_group()
        command_encoder.push_debug_group("RadixSort", (0, 0, 1))
        self.radix_sort(command_encoder,
                        self.resources.active_instance_sorting_key_buf,
                        self.resources.active_instance_primitive_list_buf,
                        self.resources.active_instance_sorting_key_sorted_buf,
                        self.resources.active_instance_primitive_list_sorted_buf,
                        self.resources.active_instance_count)
        command_encoder.pop_debug_group()
        self.encode_dispatch(command_encoder, self.resources.active_instance_count,
                             self.resources.prims_indirect_dispatch_args_buf, COMPUTE_DEFAULT_GROUP_SIZE)
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("FindTileBoundaries", (0, 0, 1))
            cursor = self.resources.bind(pass_encoder, self.FindTileBoundaries.pipeline, self.uniforms)
            pass_encoder.dispatch_compute_indirect(self.resources.prims_indirect_dispatch_args_buf)
            pass_encoder.pop_debug_group()
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("CountTileInstances", (1, 0, 1))
            cursor = self.resources.bind(pass_encoder, self.CountTileInstances.pipeline, self.uniforms)
            num_tiles = ((self.image_width + 15) // 16) *  ((self.image_height + 15) // 16)
            pass_encoder.dispatch_compute((
                (num_tiles + COMPUTE_DEFAULT_GROUP_SIZE - 1) // COMPUTE_DEFAULT_GROUP_SIZE, 1, 1
            ))
            pass_encoder.pop_debug_group()
        subres = spy.SubresourceRange()
        subres.layer, subres.layer_count, subres.mip, subres.mip_count = 0, 1, 0, 1
        command_encoder.clear_texture_float(
            texture=self.resources.output_texture,
            range=subres,
            clear_value=spy.math.float4(0, 0, 0, 0)
        )
        # Iterate over each tile. each iteration render / backpropagate gradients from 32 pixels in a tile
        for i in range(16 * 16 // 32):
            with command_encoder.begin_compute_pass() as pass_encoder:
                pass_encoder.push_debug_group("Render", (0, 1, 1))
                self.uniforms['render_tile_pixel_index_offset'] = i * 32
                cursor = self.resources.bind(pass_encoder, self.Render.pipeline, self.uniforms)
                cursor.primitive_data_buf = self.dev_primitives.blob
                cursor.d_primitive_data_buf = self.dev_primitives.d_blob
                pass_encoder.dispatch_compute((
                    (self.image_width + 15) // 16,
                    (self.image_height + 15) // 16,
                    1
                ))
                pass_encoder.pop_debug_group()
            with command_encoder.begin_compute_pass() as pass_encoder:
                pass_encoder.push_debug_group("Render_BackwardPropagation", (1, 0, 1))
                self.uniforms['render_tile_pixel_index_offset'] = i * 32
                cursor = self.resources.bind(pass_encoder, self.Render_BackwardPropagation.pipeline, self.uniforms)
                cursor.primitive_data_buf = self.dev_primitives.blob
                cursor.d_primitive_data_buf = self.dev_primitives.d_blob
                pass_encoder.dispatch_compute((
                    (self.image_width + 15) // 16,
                    (self.image_height + 15) // 16,
                    1
                ))
                pass_encoder.pop_debug_group()
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("UpdateParameters", (1, 1, 1))
            cursor = self.resources.bind(pass_encoder, self.UpdateParameters.pipeline, self.uniforms)
            cursor.primitive_data_buf = self.dev_primitives.blob
            cursor.d_primitive_data_buf = self.dev_primitives.d_blob
            cursor.ub_adam = {
                "adam_eta": self.adam_eta.value,
                "no_optimization": int(self.paused.value),
                "padding": (0, 0)
            }
            pass_encoder.dispatch_compute((
                (self.dev_primitives.blob_float_count + COMPUTE_DEFAULT_GROUP_SIZE - 1) 
                // COMPUTE_DEFAULT_GROUP_SIZE, 1, 1
            ))
            pass_encoder.pop_debug_group()
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("RegularizeParameters", (1, 1, 0))
            cursor = self.resources.bind(pass_encoder, self.RegularizeParameters.pipeline, self.uniforms)
            cursor.primitive_data_buf = self.dev_primitives.blob
            pass_encoder.dispatch_compute((
                (self.num_primitives + COMPUTE_DEFAULT_GROUP_SIZE - 1)
                // COMPUTE_DEFAULT_GROUP_SIZE, 1, 1
            ))
            pass_encoder.pop_debug_group()
    def setup_ui(self):
        screen = self.ui.screen
        window = spy.ui.Window(screen, "Settings", size=spy.float2(500, 300))

        self.fps_text = spy.ui.Text(window, "FPS: 0")
        self.loss_text = spy.ui.Text(window, "Loss: 0.0")

        def reinit_primitives():
            print('reinit primitives')
            device.wait_for_idle()
            self.reinit_primitives()

        spy.ui.Button(window, "Reinit Primitives", callback=reinit_primitives)
        self.paused = spy.ui.CheckBox(window, "Pause", False)

        self.adam_eta = spy.ui.SliderFloat(window, "Adam Eta", value=0.005, min=0.001, max=0.01)

        self.last_mouse_pos = (-1, -1)
        self.smoothed_loss_value = 0.0
    
    def on_keyboard_event(self, event):
        if self.ui.handle_keyboard_event(event):
            return
    
    def on_mouse_event(self, event : spy.MouseEvent):
        if self.ui.handle_mouse_event(event):
            return
        if event.is_move():
            if self.last_mouse_pos != (-1, -1):
                dx_pixels = event.pos.x - self.last_mouse_pos[0]
                dy_pixels = event.pos.y - self.last_mouse_pos[1]

            self.last_mouse_pos = (event.pos.x, event.pos.y)
                

if __name__ == "__main__":
    optimizer = UniformVolumePrimitivesOptimizer()
    optimizer.main()
